# -*- coding: utf-8 -*-
# file: train.py
# author: JinTian
# time: 10/05/2017 8:52 AM
# Copyright 2017 JinTian. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ------------------------------------------------------------------------
from torch.autograd import Variable

from data_loader import DataLoader, Tiny5DataLoader
from global_config import *
from models_utils import *
import torchvision
from models import fine_tune_model
from torch import nn, optim
import argparse


def exp_lr_scheduler(optimizer, epoch, init_lr=0.001, lr_decay_epoch=7):
    """Decay learning rate by a factor of 0.1 every lr_decay_epoch epochs."""
    lr = init_lr * (0.1 ** (epoch // lr_decay_epoch))

    if epoch % lr_decay_epoch == 0:
        print('LR is set to {}'.format(lr))

    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

    return optimizer


def train(args):
    data_dir = args.data_dir
    if data_dir is None or not os.path.exists(data_dir):
        print('data directory is not exist.')
    data_loader = DataLoader(data_dir=data_dir, image_size=IMAGE_SIZE, batch_size=104)
    print('data size: ', data_loader.data_sizes)
    print('data classes: ', data_loader.data_classes)
    # inputs, classes = next(iter(data_loader.load_data()))
    # print(inputs)
    # print(classes)
    # out = torchvision.utils.make_grid(inputs)
    # data_loader.show_image(out, title=[data_loader.data_classes[c] for c in classes])

    model = fine_tune_model(len(data_loader.data_classes))

    model, start_epoch = load_previous_model(model, checkpoints_dir, model_prefix=model_prefix)

    criterion = nn.CrossEntropyLoss()
    optimizer_ft = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
    epoch = 0
    try:
        since_time = time.time()
        print('training start...')

        best_acc = 0

        for epoch in range(start_epoch, num_epochs):

            optimizer = exp_lr_scheduler(optimizer_ft, epoch)
            model.train(True)  # Set model to training mode

            running_loss = 0.0
            running_corrects = 0

            for batch_data in data_loader.load_data():
                inputs, labels = batch_data
                if USE_GPU:
                    inputs, labels = Variable(inputs.cuda()), Variable(labels.cuda())
                else:
                    inputs, labels = Variable(inputs), Variable(labels)

                optimizer.zero_grad()

                outputs = model(inputs)
                _, predict = torch.max(outputs.data, 1)
                loss = criterion(outputs, labels)
                loss.backward()
                optimizer.step()

                # collect data info
                running_loss += loss.data[0]
                running_corrects += torch.sum(predict == labels.data)

            epoch_loss = running_loss / data_loader.data_sizes
            epoch_acc = running_corrects / data_loader.data_sizes

            print('Epoch: {}, {} Loss: {:.4f} Acc: {:.4f}'.format(epoch, time_since(since_time, epoch / num_epochs),
                                                                  epoch_loss, epoch_acc))
        time_elapsed = time.time() - since_time
        print('Training complete in {:.0f}m {:.0f}s'.format(
            time_elapsed // 60, time_elapsed % 60))
        print('Best val Acc: {:4f}'.format(best_acc))

        save_model(model, checkpoints_dir, model_prefix)
    except KeyboardInterrupt:
        print('manually interrupt, try saving model for now...')
        save_model(model, checkpoints_dir, model_prefix, epoch)
        print('model saved.')


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_dir', '-d', required=True, help='directory contains n classes folders inside'
                                                                'are images.')

    args = parser.parse_args()
    return args


def main():
    args = parse_args()

    train(args)


if __name__ == '__main__':
    main()
